import pandas as pd
import numpy as np
import tushare as ts
import math
import operator
from scipy.spatial.distance import euclidean
from fastdtw import fastdtw

# 获得距离
def dtw_dist(x,y):
    distance,path = fastdtw(x,y,dist=euclidean)
    return distance

#获得相邻元素
def getN(training,testing,k,n):
    distances = []
    for x in range(n):
        dist = dtw_dist(testing,training[x])
        distances.append((training[x],dist))
    distances.sort(key=operator.itemgetter(1))
    neighbors = []
    for x in range(k):
        neighbors.append(distances[x][0])
    return neighbors

#对选定元素进行投票，获得最终预测
def getR(neighbors):
    classvotes = {}
    for x in range(len(neighbors)):
        response = neighbors[x][-1]
        if response in classvotes:
            classvotes[response] += 1
        else:
            classvotes[response] = 1
    sortedvotes = sorted(classvotes.items(),key=operator.itemgetter(1),reverse=True)
    return sortedvotes[0][0]

#计算准确率
def getA(testing,predictions):
    correct = 0
    for x in range(len(testing)):
        if testing[x][-1] == predictions[x]:
            correct += 1
    return (correct/float(len(testing))) * 100.0

#主函数
def main(code,k,ktype='D'):
#获得股票数据并写入多维数组
    sp = ts.get_k_data(str(code),ktype=str(ktype),start='2017-01-01')

    foreseecl = sp.assign(foresee = 0)  

    for i in range(len(sp)-1):
        j = i + 1
        
        dclose = sp.get_value(j,'close')
        dopen = sp.get_value(j,'open')
        if dclose - dopen > 0:
            predict = foreseecl.set_value(i,'foresee',1)
        else:
            predict = foreseecl.set_value(i,'foresee',0)
    #print(predict)

#构建训练集和测试集，其中将要进行预测的数据写入测试集
    testing = []
    training = []
    for j in range(k):
        
        training.append([]) 
        d_open = foreseecl.get_value(j,'open')
        d_close = foreseecl.get_value(j,'close')
        d_high = foreseecl.get_value(j,'high')
        d_low = foreseecl.get_value(j,'low')
        d_volume = foreseecl.get_value(j,'volume')
        d_foresee = foreseecl.get_value(j,'foresee')

        k1 = (d_close - d_open)/d_open
        k2 = (d_high - d_low)/(d_high + d_low)
        k3 = d_volume 

        training[-1].append(k1)
        training[-1].append(k2)
        training[-1].append(k3)
        training[-1].append(d_foresee)

    for j in range(len(predict)-1-k):
        
        training.append([])
        testing.append([])  
        d_open = foreseecl.get_value(j,'open')
        d_close = foreseecl.get_value(j,'close')
        d_high = foreseecl.get_value(j,'high')
        d_low = foreseecl.get_value(j,'low')
        d_volume = foreseecl.get_value(j,'volume')
        d_foresee = foreseecl.get_value(j,'foresee')

        k1 = (d_close - d_open)/d_open
        k2 = (d_high - d_low)/(d_high + d_low)
        k3 = d_volume 

        training[-1].append(k1)
        training[-1].append(k2)
        training[-1].append(k3)
        training[-1].append(d_foresee)

        testing[-1].append(k1)
        testing[-1].append(k2)
        testing[-1].append(k3)
        testing[-1].append(d_foresee)

    x = len(predict) - 1
    testing.append([]) 
    d_open = foreseecl.get_value(x,'open')
    d_close = foreseecl.get_value(x,'close')
    d_high = foreseecl.get_value(x,'high')
    d_low = foreseecl.get_value(x,'low')
    d_volume = foreseecl.get_value(x,'volume')
    d_foresee = foreseecl.get_value(x,'foresee')

    k1 = (d_close - d_open)/d_open
    k2 = (d_high - d_low)/(d_high + d_low)
    k3 = d_volume 

    testing[-1].append(k1)
    testing[-1].append(k2)
    testing[-1].append(k3)
    testing[-1].append(d_foresee)
    

#对每一组数据进行预测汇总并测试准确率
    predictions=[]
    
    for x in range(len(testing)):
        t = x + k
        neighbors = getN(training,testing[x],k,t)
        result = getR(neighbors)
        predictions.append(result)
        #print('> predicted=' + repr(result) + ', actual=' + repr(testing[x][-1]))
    accuracy = getA(testing, predictions)
    print('k = ' + str(k))
    print('code = ' + str(code))
    print('Accuracy: ' + repr(accuracy) + '%')
    if predictions[-1] == 1:
        print('buy in')
    else:
        print('sell out') 

    return accuracy, predictions[-1]


    







